from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_sslify import SSLify

import random
import torch
import json
import math
import numpy as np
import os

import a2c_FoV as a2c
import Arguments
from prepareFoVdata import *
from torch.distributions import Categorical

app = Flask(__name__)
# sslify = SSLify(app)
CORS(app, supports_credentials=True, resource=r'/*')
server_cert = './certificate/server/server-cert.pem'
server_key = './certificate/server/server-key.pem'

# ABR paremeters
S_INFO = 6
# bit_rate, buffer_size, rebuffering_time, bandwidth_measurement, chunk_til_video_end
S_LEN = 8  # take how many frames in the past
A_DIM = 4  #
ACTOR_LR_RATE = 0.0001
BUFFER_NORM_FACTOR = 10.0  # unit：s
CHUNK_TIL_VIDEO_END_CAP = 200
RAND_RANGE = 1000
M_IN_K = 1000.0
DEVIATION_THRESHOLD = 0.25
NN_MODEL = './model/actor_nn_model_ep_640.pkl'

# FoV parameters
videoName = "tiledSkiing"
videoId = 0
isLoad = False
sal_maps = None
time_array = None
FoV_net = convlstm.ConvLSTM_model(input_dim=args.input_size,
                                  hidden_dim=args.hidden_size,
                                  kernel_size=(5, 5),
                                  num_layers=args.num_layers,
                                  batch_first=True)


def convert_torch(variable, dtype=np.float32):
    if variable.dtype != dtype:
        variable = variable.astype(dtype)
    return torch.from_numpy(variable)


@app.route('/ABR', methods=['POST'])
def ABR():
    face_idx = 0
    tile_idx = 0
    try:
        coord = request.json['coord']
    except TypeError:
        print("Get data some error osh")
        return '0'
    # print(coord)
    buffer = request.json['buffer']
    if buffer is None:
        buffer = 0
    # print(buffer)
    lastQuality = request.json['lastQuality']
    # print(lastQuality)
    past_8_index = request.json['past_8_index']
    # print(past_8_index)
    past_8_throughput = request.json['past_8_throughput']
    # print(past_8_throughput)
    tile_dev = request.json['tile_dev']
    # movie_name = request.json['movie_name']
    movie_name = 'tiledSkiing'
    # print(movie_name)
    index = request.json['index']
    # print(index)
    face = request.json['face']
    # print(face)

    # 加载电影json数据
    with open('./model/tiledSkiing.json', 'r') as load_f:
        movie_json = json.load(load_f)

    bitrates = movie_json['bitrate_bps']
    if movie_name == 'tiledSkiing':
        face_idx = face // 4
        tile_idx = face % 4
        # FACE_COORD = [[1, 0.5, -0.5], [1, 0.5, 0.5], [1, -0.5, -0.5], [1, -0.5, 0.5],
        #               [-1, 0.5, 0.5], [-1, 0.5, -0.5], [-1, -0.5, 0.5], [-1, -0.5, -0.5],
        #               [-0.5, 1, 0.5], [0.5, 1, 0.5], [-0.5, 1, -0.5], [0.5, 1, -0.5],
        #               [-0.5, -1, -0.5], [0.5, -1, -0.5], [-0.5, -1, 0.5], [0.5, -1, 0.5],
        #               [-0.5, 0.5, -1], [0.5, 0.5, -1], [-0.5, -0.5, -1], [0.5, -0.5, -1],
        #               [0.5, 0.5, 1], [-0.5, 0.5, 1], [0.5, -0.5, 1], [-0.5, -0.5, 1]]
    # else:
        # FACE_COORD = [[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, -1], [0, 0, 1]]

    # face_coord = FACE_COORD[face] deviation = math.acos((coord[0] * face_coord[0] + coord[1] * face_coord[1] +
    # coord[2] * face_coord[2]) / math.sqrt(coord[0] * coord[0] + coord[1] * coord[1] + coord[2] * coord[2]) /
    # math.sqrt(face_coord[0] * face_coord[0] + face_coord[1] * face_coord[1] + face_coord[2] * face_coord[2]))
    # deviation /= math.pi
    fov_tmp = 0
    # if deviation < DEVIATION_THRESHOLD:
    if tile_dev <= 20:
        fov_tmp = 1

    past_download_time = np.zeros(8)
    for i in range(8):
        if past_8_throughput[i] == 0 or past_8_throughput[i] is None:
            past_8_throughput[i] = 0
            past_download_time[i] = 0
        else:
            if movie_name == 'tiledSkiing':
                past_download_time[i] = movie_json['segment_size_bytes'][max(0, index - 8 + i)][face_idx][tile_idx][
                                            past_8_index[i]] / past_8_throughput[i]
            else:
                past_download_time[i] = movie_json['segment_size_bytes'][max(0, index - 8 + i)][face][past_8_index[i]] / \
                                        past_8_throughput[i]
            past_8_throughput[i] /= M_IN_K
            past_8_throughput[i] /= M_IN_K
            past_8_throughput[i] /= M_IN_K

    movie_segment_size = np.zeros(4)
    if movie_name == 'tiledSkiing':
        movie_segment_size[0] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][0] / M_IN_K / M_IN_K
        movie_segment_size[1] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][1] / M_IN_K / M_IN_K
        movie_segment_size[2] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][2] / M_IN_K / M_IN_K
        movie_segment_size[3] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][3] / M_IN_K / M_IN_K
    else:
        movie_segment_size[0] = movie_json['segment_size_bytes'][index][face][0] / M_IN_K / M_IN_K
        movie_segment_size[1] = movie_json['segment_size_bytes'][index][face][1] / M_IN_K / M_IN_K
        movie_segment_size[2] = movie_json['segment_size_bytes'][index][face][2] / M_IN_K / M_IN_K
        movie_segment_size[3] = movie_json['segment_size_bytes'][index][face][3] / M_IN_K / M_IN_K

    np.random.seed(20)
    actor = a2c.ActorNet(s_dim=[S_INFO, S_LEN], a_dim=A_DIM, lr=ACTOR_LR_RATE)
    actor.load_state_dict(torch.load(NN_MODEL))
    state = np.zeros((S_INFO, S_LEN))

    # 上一个视频块比特率
    state[0, 7] = fov_tmp
    # 当前缓冲大小
    state[1, :] = past_8_throughput  # KBps
    # 过去8个视频块对应网络吞吐量
    state[2, :] = past_download_time  # 10 sec
    # 过去8个视频块下载时间
    state[3, :A_DIM] = movie_segment_size
    # 未来视频块不同质量大小
    state[4, :] = buffer / BUFFER_NORM_FACTOR
    # 剩余视频块数量
    state[5, 7] = bitrates[lastQuality] / float(np.max(bitrates))
    # print(state)

    # 随机性选择动作
    _, _, action_prob = actor.get_actor_out(convert_torch(np.reshape(state, (1, S_INFO, S_LEN))))
    # 随机策略
    # action_prob = action_prob.cpu()
    # m = Categorical(action_prob)
    # bit_rate = m.sample().detach().item()
    # 确定性策略
    action_prob = action_prob.numpy()
    bit_rate = action_prob.argmax()

    print(coord, tile_dev, action_prob, bit_rate)

    return str(bit_rate)


@app.route('/FoV', methods=['POST'])
def FoV():
    start_t = time.time()
    req = request.get_json(force=True)
    # print('past FoV:', type(req))
    req.reverse()
    # for item in req:
    #     print(item)
    if not req or len(req) <= 3:
        return "the length of post_data < 4"
    inputs, pre_time = get_sal_fix(time_array, sal_maps, req)

    # return jsonify([pre_time, [-0.7,0,-0.7]])
    # output_list 每个时间点最后一层隐层结果: list
    # output_last 最后一个时间点隐层和细胞层结果: list
    # conv_output 最后一个时间点隐层和细胞层卷积之后结果: list
    # predict_array 每个时间点隐层卷积之后结果 torch.tensor
    output_list, output_last, conv_output, conv_output_list_ret = FoV_net(inputs)
    pre_fov, pre_tiles = return_fov(conv_output)
    end_t = time.time()
    print(f"Time={round(end_t - start_t, 4)}")
    return jsonify([pre_time, pre_fov])


@app.route('/', methods=['GET'])
def home():
    global videoName, videoId, isLoad, sal_maps, time_array, FoV_net

    videoName = request.args.get("videoName")
    if videoName is None:
        return "Backend of the server! None parameter!"
    videoName = 'Skiing'

    sal_maps = sal_maps_list[videoId]
    time_array = time_array_list[videoId]
    print(f'videoName={videoName}, videoId={videoId}, '
          f'sal_maps\' shape={sal_maps.shape}, time_array\'s shape={time_array.shape}')
    if isLoad:
        return f"Success! Have Loaded model parameters!"

    if torch.cuda.is_available():
        FoV_net = FoV_net.cuda()
    if not os.path.exists(args.model_path):
        exit("model path doesn't exist")
    # modelPath = args.model_path + f"convlstm_offline_{videoName.lower()}_1_1s.pth"
    modelPath = args.model_path + f"convlstm_offline_{videoName.lower()}_1_4s.pth"
    if not os.path.exists(modelPath):
        exit(f"{modelPath} doesn't exit!")

    FoV_net.load_state_dict(torch.load(modelPath))
    print('Loading model from', modelPath)
    print("Total number of parameters in networks is {}".format(sum(x.numel() for x in FoV_net.parameters())))
    FoV_net.eval()
    isLoad = True
    return f"Success! Have gotten videoName={videoName}!"


if __name__ == '__main__':
    # app.run(host='10.112.212.127', port=9009, debug=True, ssl_context=(server_cert, server_key))
    # app.run(host='10.112.55.254', port=9009, debug=True, ssl_context=(server_cert, server_key))
    app.run(host='0.0.0.0', port=9009, debug=True, threaded=False)
    # app.run(host='192.168.0.197', port=9009, debug=True, ssl_context=(server_cert, server_key))
    # app.run(host='0.0.0.0', port=9009, debug=True, ssl_context=(server_cert, server_key))

